#include <iostream>
#include <iomanip>
#include <sstream>
#include <cstdlib>
#include <map>
#include <algorithm>
#include <fstream>
#include <time.h>

#include "fastjet/ClusterSequence.hh"
#include "fastjet/tools/Pruner.hh"

#include "Njettiness.hh"

#include "CmdLine.hh"
#include "TFile.h"
#include "TChain.h"
#include "TChainElement.h"

#include "TH1F.h"
#include "TH2F.h"

#include "Analysis.h"
#include "fastjet/tools/Filter.hh"


class FileSkimmer : public Analysis
{
public:

  const static int nmax=100;

  Double_t tau1_1[100];
  Double_t tau2_1[100];
  Double_t tau3_1[100];
  Double_t tau4_1[100];

  Double_t tau1_2[100];
  Double_t tau2_2[100];
  Double_t tau3_2[100];
  Double_t tau4_2[100];

  Double_t tau1_kt[100];
  Double_t tau2_kt[100];
  Double_t tau3_kt[100];
  Double_t tau4_kt[100];

  Double_t tau1_3[100];
  Double_t tau2_3[100];
  Double_t tau3_3[100];
  Double_t tau4_3[100];
  
  Double_t jet_pt[100];
  Double_t jet_m[100];
  
  Double_t jet_px[100];
  Double_t jet_py[100];
  Double_t jet_pz[100];
  Double_t jet_e[100];

  Int_t jet_il[100];
  Int_t jet_ih[100];
  

  
  Double_t top_pt[100];
  Int_t top_size;

  //normalization of the taus
  Double_t d[100];

  Double_t met_px, met_py;

  Int_t jet_size;

  

  //skinny jets
  Double_t sjet_pt[100];
  Double_t sjet_m[100];
  
  Double_t sjet_px[100];
  Double_t sjet_py[100];
  Double_t sjet_pz[100];
  Double_t sjet_e[100];

  Int_t sjet_size;


  //subjets
  Double_t subjet_pt[200];
  
  Double_t subjet_px[200];
  Double_t subjet_py[200];
  Double_t subjet_pz[200];
  Double_t subjet_e[200];

  Int_t subjet_size;

  

  Double_t MJ;


  TTree* treeout;
  FileSkimmer(const CmdLine& cmdline)
    : Analysis(cmdline)
  {
    fileout->cd();
    treeout=new TTree("Ntuple", "Ntuple");

    hist["met"]=new TH1D("MEt", "#slash{E}_{T}", 50, 0, 500);

    for(int i=0; i<20; i++)
      {
	stringstream sout;
	sout<<"skinny_jet"<<i+1;
	stringstream sout1;
	sout1<<"Skinny Jet "<<i+1;
	hist[sout.str()]=new TH2D(sout.str().c_str(),
				  sout1.str().c_str(),
				  64, -3.1914592, 3.1914592,
				  63, 0, 3.1415926*2);
	hist[sout.str()]->SetOption("colz");
	
	if(i<6)
	  {
	    stringstream tout;
	    tout<<"fat_jet"<<i+1;
	    stringstream tout1;
	    tout1<<"Fat Jet "<<i+1;
	    hist[tout.str()]=new TH2D(tout.str().c_str(),
				      tout1.str().c_str(),
				      64, -3.1914592, 3.1914592,
				      63, 0, 3.1415926*2);
	    hist[tout.str()]->SetOption("colz");
	  }

      }


    treeout->Branch("weight", &weight, "weight/D");  

    treeout->Branch("n", &jet_size, "n/I");
    treeout->Branch("e", jet_e, "e[n]/D");
    treeout->Branch("px", jet_px, "px[n]/D");
    treeout->Branch("py", jet_py, "py[n]/D");
    treeout->Branch("pz", jet_pz, "pz[n]/D");
    treeout->Branch("pt", jet_pt, "pt[n]/D");
    treeout->Branch("m", jet_m, "m[n]/D");
    treeout->Branch("il", jet_il, "il[n]/I");
    treeout->Branch("ih", jet_ih, "ih[n]/I");
    treeout->Branch("d", d, "d[n]/D");
    


    treeout->Branch("met_px", &met_px, "met_px/D");
    treeout->Branch("met_py", &met_py, "met_py/D");

    
    treeout->Branch("s_n", &sjet_size, "s_n/I");
    treeout->Branch("s_e", sjet_e, "s_e[s_n]/D");
    treeout->Branch("s_px", sjet_px, "s_px[s_n]/D");
    treeout->Branch("s_py", sjet_py, "s_py[s_n]/D");
    treeout->Branch("s_pz", sjet_pz, "s_pz[s_n]/D");
    treeout->Branch("s_pt", sjet_pt, "s_pt[s_n]/D");
    treeout->Branch("s_m", sjet_m, "s_m[s_n]/D");
    
    treeout->Branch("sub_n", &subjet_size, "sub_n/I");
    treeout->Branch("sub_e", subjet_e, "sub_e[sub_n]/D");
    treeout->Branch("sub_px", subjet_px, "sub_px[sub_n]/D");
    treeout->Branch("sub_py", subjet_py, "sub_py[sub_n]/D");
    treeout->Branch("sub_pz", subjet_pz, "sub_pz[sub_n]/D");
    treeout->Branch("sub_pt", subjet_pt, "sub_pt[sub_n]/D");
    

    treeout->Branch("tau1_1", tau1_1, "tau1_1[n]/D");
    treeout->Branch("tau2_1", tau2_1, "tau2_1[n]/D");
    treeout->Branch("tau3_1", tau3_1, "tau3_1[n]/D");
    treeout->Branch("tau4_1", tau4_1, "tau4_1[n]/D");

    treeout->Branch("tau1_2", tau1_2, "tau1_2[n]/D");
    treeout->Branch("tau2_2", tau2_2, "tau2_2[n]/D");
    treeout->Branch("tau3_2", tau3_2, "tau3_2[n]/D");
    treeout->Branch("tau4_2", tau4_2, "tau4_2[n]/D");

    treeout->Branch("tau1_kt", tau1_kt, "tau1_kt[n]/D");
    treeout->Branch("tau2_kt", tau2_kt, "tau2_kt[n]/D");
    treeout->Branch("tau3_kt", tau3_kt, "tau3_kt[n]/D");
    treeout->Branch("tau4_kt", tau4_kt, "tau4_kt[n]/D");
    

    treeout->Branch("top_n", &top_size, "top_n/I");
    treeout->Branch("top_pt", top_pt, "top_pt[top_n]/D");

    /*
    treeout->Branch("tau1_3", tau1_3, "tau1_3[n]/D");
    treeout->Branch("tau2_3", tau2_3, "tau2_3[n]/D");
    treeout->Branch("tau3_3", tau3_3, "tau3_3[n]/D");
    */
    
  }

  bool cut()
  {
    hist["nevt"]->Fill(0.0,weight);
    hist["nevt_unweighted"]->Fill(0.0);


    if(cmdline.present("-Mt"))
      {
	float Mt_cut=cmdline.value<float>("-Mt", 400);
	float Mt_present=0;
	for(int i=0; i<jets.size(); i++)
	  if(jets[i].pt()>50.0)
	  {Mt_present+=jets[i].m();}
	if(Mt_present < Mt_cut)
	  return false;
      }

    if(cmdline.present("-jet4"))
      {
	if(Event_==0)
	  {
	    cout<<"OPTIONS: Demand 1j > 100, at least 4j > 50"<<endl;
	    hist["nevt"]->GetXaxis()->SetBinLabel(2, "jet4");
	  }
	
	if(jets.size()<4)
	  return false;
	if(jets[0].pt()<100)
	  return false;
	if(jets[3].pt() <50.0)
	  return false;
      }

    
    hist["nevt"]->Fill(1.0,weight);
    hist["nevt_unweighted"]->Fill(1.0);
    return true;
  }

  void analyze();

};



int main(int argc, char** argv)
{

  //initialize command line interface
  CmdLine cmdline(argc,argv);

  if(!cmdline.present("-i") || !cmdline.present("-o"))
    {
      cout<<"Ntuplizer: trim down rootfile by applying predefined cuts";
      cout<<"Usage: -i <input> -o <output> -a [algorithm] -r [radius] -m [mass] -n [nev] -p -w"<<endl;
      cout<<"-a \t clustering algorithm to be used"<<endl;
      cout<<"-r \t clustering radius"<<endl;
      cout<<"-n \t number of events to process"<<endl;
      cout<<"-p \t whether jet pruning is used"<<endl;
      cout<<"-w \t whether event has a weight stored as evt Number"<<endl;
      cout<<"-gluino \t retrieve plots that trace the gluinos "<<endl;
      cout<<"-jet4 \t apply 4 jets cut "<<endl;
      cout<<"-jet2 \t apply 2 jets cut "<<endl;
      cout<<"-display \t number of display plots"<<endl;
      
      return 1;
    }

  //create analysis object
  FileSkimmer myanalysis(cmdline);
  myanalysis.tree_init();
  myanalysis.analyze();

  return 0;
}


//main function to loop over all the entries
void FileSkimmer::analyze()
{
  //setup N-subjettiness

  Njettiness::AxesMode mode=Njettiness::min_axes;

  /*
  if(cmdline.present("-min-axes"))
    {
      mode=Njettiness::min_axes;
      cout<<"MESSAGE: minimial axes mode enabled (slow)"<<endl;
    }
  */

  NsubParameters param1(1.0, 1.2, 100);
  Njettiness nSubMin1(mode,
		      param1);
  
  NsubParameters param2(2.0, 1.2, 100);
  Njettiness nSubMin2(mode,
		      param2);
  
  NsubParameters paramkt(1.0, 1.2, 100);
  Njettiness nSubMinkt(mode,
		       paramkt);
  
  NsubParameters param3(3, 1.2, 100);
  Njettiness nSubMin3(mode,
		      param3);
  
  //initialize trimmer
  Filter trimmer(0.3, SelectorPtFractionMin(0.05));

  /*
  NsubParameters param1(1.0, 1.2, 100);
  Njettiness nSubMin1(Njettiness::min_axes,//min_axes,
		      param1);
  
  NsubParameters param2(2.0, 1.2, 100);
  Njettiness nSubMin2(Njettiness::min_axes,
		      param2);
  
  NsubParameters paramkt(1.0, 1.2, 100);
  Njettiness nSubMinkt(Njettiness::kt_axes,
		       paramkt);
  
  NsubParameters param3(3, 1.2, 100);
  Njettiness nSubMin3(Njettiness::min_axes,
		      param3);
  */


  //begin event loop
  int entry=skipevt; 

  for(;(maxevt<0 || entry<maxevt+skipevt) && entry<tree.GetEntries();
      entry++)
    {
      
      hist["nevt"]->Fill(0.0,weight);
      hist["nevt_unweighted"]->Fill(0.0);

      if(entry%1000==0 && entry/1000 > 0)
	cout<<"MESSAGE: "<<entry<<" Processed"<<endl;
      
      //populate our variables
      tree.GetEntry(entry);
      jet_size=0;
      sjet_size=0;
      subjet_size=0;

      Event_=entry;
      if(cmdline.present("-unit-weight"))
	{
	  weight=1.0;
	}
      
      if(weight==0.0)
	weight=1.0;


      //make skinny jets first
      makejets(antikt_algorithm, 0.4, 10.0);
      
      if(jets.size()<2) continue;

      top_size=0;
      //grab top pt
      for(int tt=0; tt<ttop_pt.size(); tt++)
	{
	  top_pt[tt]=ttop_pt[tt];
	  top_size++;
	}

      //grab the skinny jets
      sjet_size=jets.size();
      for(int i=0; i<jets.size(); i++)
	{
	  sjet_px[i]=jets[i].px();
	  sjet_py[i]=jets[i].py();
	  sjet_pz[i]=jets[i].pz();
	  sjet_pt[i]=jets[i].pt();
	  sjet_e[i]=jets[i].e();
	  sjet_m[i]=jets[i].m();

	  vector<PseudoJet> sub=jets[i].constituents();
	  stringstream sout;
	  sout<<"skinny_jet"<<i+1;

	  for(int k=0; k<sub.size(); k++)
	    {
	      ((TH2*)(hist[sout.str()].p))->Fill(sub[k].eta(), sub[k].phi(),
						  sub[k].pt());
	    }
	}

      

      //make fat-jets
      makejets();
      
      //trim all the jets
      if(b_trim)
	for(int i=0; i<jets.size(); i++)
	  jets[i] = trimmer(jets[i]);
	

      //resort the jets after trimming
      jets=sorted_by_pt(jets);
      
      //get the MEt
      met_px=met.px();
      met_py=met.py();
      
      cout<<endl;
      cout<<"Event: "<<entry<<endl;
      

      double MJ=0;
      double HT=0;
      int Tcount=0;
      double T21=1, T32=1, T43=1;

      for(int i=0; i<jets.size(); i++)
	{
	  double c_pt = jets[i].pt();
	  
	  if(c_pt>50.0)
	    {
	      MJ+=jets[i].m();
	      HT+=sqrt(jets[i].m()*jets[i].m()+c_pt*c_pt);
	    }

	  vector<PseudoJet> sub=jets[i].constituents();

	  stringstream sout;
	  sout<<"fat_jet"<<i+1;

	  for(int k=0; k<sub.size(); k++)
	    {
	      ((TH2*)(hist[sout.str()].p))->Fill(sub[k].eta(), sub[k].phi(),
						 sub[k].pt());
	    }

	  
	  cout<<"jet: "<<i<<" size: "<<sub.size()<<endl;

	  jet_e[jet_size]=jets[i].e();
	  jet_px[jet_size]=jets[i].px();
	  jet_py[jet_size]=jets[i].py();
	  jet_pz[jet_size]=jets[i].pz();
	  jet_pt[jet_size]=jets[i].pt();
	  jet_m[jet_size]=jets[i].m();
	  jet_il[jet_size]=subjet_size;
	  jet_ih[jet_size]=-1;

	  //grab the d variable
	  d[jet_size]=0;
	  
	  for(int k=0; k<sub.size(); k++)
	    d[jet_size]+=sub[k].pt();
	  
	  d[jet_size]*=1.2;

	  tau1_1[jet_size]= nSubMin1.getTau(1, sub);
	  tau2_1[jet_size]= nSubMin1.getTau(2, sub);
	  tau3_1[jet_size]= nSubMin1.getTau(3, sub);
	  tau4_1[jet_size]= nSubMin1.getTau(4, sub);
	  
	  tau1_2[jet_size]= nSubMin2.getTau(1, sub);
	  tau2_2[jet_size]= nSubMin2.getTau(2, sub);
	  tau3_2[jet_size]= nSubMin2.getTau(3, sub);
	  tau4_2[jet_size]= nSubMin2.getTau(4, sub);

	  if(Tcount<4)
	    {
	      T21*=tau2_1[jet_size]/tau1_1[jet_size];
	      T32*=tau3_1[jet_size]/tau2_1[jet_size];
	      T43*=tau4_1[jet_size]/tau3_1[jet_size];
	      Tcount++;
	    }

	  /*
	  tau1_3[jet_size]= nSubMin3.getTau(1, sub);
	  tau2_3[jet_size]= nSubMin3.getTau(2, sub);
	  tau3_3[jet_size]= nSubMin3.getTau(3, sub);
	  */


	  tau1_kt[jet_size]= nSubMinkt.getTau(1, sub);
	  tau2_kt[jet_size]= nSubMinkt.getTau(2, sub);
	  tau3_kt[jet_size]= nSubMinkt.getTau(3, sub);
	  tau4_kt[jet_size]= nSubMinkt.getTau(4, sub);
	  
	  vector<PseudoJet> subjets=jets[i].pieces();
	  for(int j=0; j<subjets.size(); j++)
	    {
	      subjet_px[subjet_size]=subjets[j].px();
	      subjet_py[subjet_size]=subjets[j].py();
	      subjet_pz[subjet_size]=subjets[j].pz();
	      subjet_e[subjet_size]=subjets[j].e();
	      subjet_pt[subjet_size]=subjets[j].pt();
	      subjet_size++;
	    }
	  jet_ih[jet_size]=subjet_size;

	  cout<<jet_size<<": "<<jets[i].m()
	  <<","<<jets[i].pt()
	  <<","<<jets[i].eta()
	  <<","<<jets[i].phi()<<endl;

	  cout<<"tau: "<<tau2_1[jet_size]/tau1_1[jet_size]
	      <<","<<tau3_1[jet_size]/tau2_1[jet_size]
	      <<","<<tau4_1[jet_size]/tau3_1[jet_size]<<endl;
	  
	  jet_size++;
	}
      cout<<"MJ: "<<MJ<<" HT: "<<HT<< " Ts: "<<pow(T21,0.25)<<","<<pow(T32,0.25)<<","<<pow(T43,0.25)<<endl;

      /*      
      makejets(antikt_algorithm, 0.7);

      for(int i=0; i<jets.size(); i++)
	{
	  if(jets[i].pt()<10.0)
	    break;

	  sjet_e[sjet_size]=jets[i].e();
	  sjet_px[sjet_size]=jets[i].px();
	  sjet_py[sjet_size]=jets[i].py();
	  sjet_pz[sjet_size]=jets[i].pz();
	  sjet_pt[sjet_size]=jets[i].pt();
	  sjet_m[sjet_size]=jets[i].m();
	  
	  sjet_size++;
	}
      */

      

      treeout->Fill();
      
    }


  treeout->Write();

  cout<<"MESSAGE: total entry processed: "
      <<entry<<endl;

  cout<<"MESSAGE: total weight processed: "
      <<hist["nevt"]->GetBinContent(1)<<endl;




  //loop over all the TFiles
  //grab the precut # of events

  //already implemented in Analysis.h

  /*
  TObjArray *fileElements=tree.GetListOfFiles();
  TIter next(fileElements);
  TChainElement *chEl=0;

 
  int c_new_nevt=0;
  int c_new_nevt_unweighted=0;
  bool b_fix_new_nevt=false;

  while (( chEl=(TChainElement*)next() )) {
    TFile f(chEl->GetTitle());
    TH1D* c_nevt=(TH1D*) f.Get("nevt");
    TH1D* c_nevt_unweighted=(TH1D*) f.Get("nevt_unweighted");
    if(c_nevt==NULL || c_nevt_unweighted==NULL)
      continue;
    
    b_fix_new_nevt=true;
    c_new_nevt+=c_nevt->GetBinContent(1);
    c_new_nevt_unweighted+=c_nevt_unweighted->GetBinContent(1);
  }


  if(b_fix_new_nevt)
    {
      hist["nevt"]->Fill(0.0,c_new_nevt);
      hist["nevt_unweighted"]->Fill(0.0,c_new_nevt_unweighted);
      cout<<"MESSAGE: pre-processed sum of weights: "<<c_new_nevt<<endl;
    }
  */


  
  fileout->Write();
  fileout->Close();
}


